// SPDX-License-Identifier: GPL-2.0
#include <linux/types.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/sched.h>
#include <linux/stddef.h>
#include <linux/slab.h>
#include <linux/proc_fs.h>
#include <linux/rculist.h>
#include <linux/rhashtable.h>
#include <linux/netfilter.h>
#include <net/netfilter/nf_conntrack.h>
#include <net/net_namespace.h>

struct nf_conn_statis_key {
	u16 l3num;
	union {
		u32 ipv4;
		struct in6_addr ipv6;
	};
};

struct nf_conn_statis_entry {
	struct rhash_head rhnode;
	struct nf_conn_statis_key key;
	struct hlist_node hlnode;
	u32 is_report;
	unsigned long report_expires;
	unsigned long update_expires;
	atomic_t per_sec_use; /* nfct nums per second */
	atomic_t use; /* total nfct nums */
	struct rcu_head rcu;
};

static const struct rhashtable_params nf_conn_rht_params = {
	.head_offset = offsetof(struct nf_conn_statis_entry, rhnode),
	.key_offset = offsetof(struct nf_conn_statis_entry, key),
	.key_len = sizeof(struct nf_conn_statis_key),
	.automatic_shrinking = true,
};

static struct kmem_cache *g_nf_conn_statis_cache __read_mostly;
static struct rhashtable g_nf_conn_statis_tbl;
static struct hlist_head g_nf_conn_statits_list;
static atomic_t g_nf_conn_statis_num;
static spinlock_t g_nf_conn_statis_lock;

static unsigned int g_nf_conn_statis_single_max __read_mostly = 1000;
module_param(g_nf_conn_statis_single_max, uint, 0600);
MODULE_PARM_DESC(g_nf_conn_statis_single_max, "g_nf_conn_statis_max_num");

static unsigned int g_nf_conn_statis_max_num __read_mostly = 1000;
module_param(g_nf_conn_statis_max_num, uint, 0600);
MODULE_PARM_DESC(g_nf_conn_statis_max_num, "g_nf_conn_statis_max_num");

static unsigned int g_nf_conn_statis_switch __read_mostly;
module_param(g_nf_conn_statis_switch, uint, 0600);
MODULE_PARM_DESC(g_nf_conn_statis_switch, "g_nf_conn_statis_switch");

void (*g_nf_conn_statis_report)(void *key, u32 num_per_sec);
EXPORT_SYMBOL(g_nf_conn_statis_report);

static inline struct nf_conn_statis_entry *
nf_conn_statis_find(struct rhashtable *tbl, struct nf_conn_statis_key *key)
{
	struct nf_conn_statis_entry *nf_conn;

	rcu_read_lock();
	nf_conn = rhashtable_lookup(tbl, key, nf_conn_rht_params);
	rcu_read_unlock();

	return nf_conn;
}

static inline bool
nf_conn_statis_get_key(const struct nf_conn *ct, struct nf_conn_statis_key *key)
{
	const struct nf_conntrack_tuple *tuple;

	if (!g_nf_conn_statis_switch)
		return false;

	tuple = &ct->tuplehash[IP_CT_DIR_ORIGINAL].tuple;

	key->l3num = tuple->src.l3num;
	if (tuple->src.l3num == NFPROTO_IPV4)
		key->ipv4 = tuple->src.u3.ip;
	else if (tuple->src.l3num == NFPROTO_IPV6)
		key->ipv6 = tuple->src.u3.in6;
	else
		return false;

	return true;
}

static inline void nf_conn_statis_rcu_free(struct rcu_head *head)
{
	struct nf_conn_statis_entry *nf_conn;

	nf_conn = container_of(head, struct nf_conn_statis_entry, rcu);
	kmem_cache_free(g_nf_conn_statis_cache, nf_conn);
}

void nf_conn_statis_try_insert_entry(const struct nf_conn *ct)
{
	struct nf_conn_statis_entry *nf_conn;
	struct nf_conn_statis_key key = {0};

	if (!nf_conn_statis_get_key(ct, &key))
		return;

	spin_lock_bh(&g_nf_conn_statis_lock);
	nf_conn = nf_conn_statis_find(&g_nf_conn_statis_tbl, &key);
	if (nf_conn) {
		atomic_inc(&nf_conn->use);
		if (time_after_eq(jiffies, nf_conn->update_expires)) {
			nf_conn->update_expires = jiffies + HZ;
			atomic_set(&nf_conn->per_sec_use, 1);
			goto out;
		} else {
			atomic_inc(&nf_conn->per_sec_use);
		}

		if (!nf_conn->is_report &&
		    (atomic_read(&nf_conn->per_sec_use) > g_nf_conn_statis_single_max)) {
			if (g_nf_conn_statis_report)
				g_nf_conn_statis_report(&key, atomic_read(&nf_conn->per_sec_use));
			nf_conn->is_report = 1;
			nf_conn->report_expires = jiffies;
		}
		goto out;
	}

	if (atomic_read(&g_nf_conn_statis_num) > g_nf_conn_statis_max_num)
		goto out;

	nf_conn = kmem_cache_alloc(g_nf_conn_statis_cache, GFP_ATOMIC);
	if (!nf_conn)
		goto out;

	nf_conn->key = key;
	nf_conn->is_report = 0;
	atomic_set(&nf_conn->use, 1);
	atomic_set(&nf_conn->per_sec_use, 1);
	nf_conn->update_expires = jiffies + HZ;
	if (rhashtable_lookup_insert_fast(&g_nf_conn_statis_tbl, &nf_conn->rhnode,
					  nf_conn_rht_params)) {
		kmem_cache_free(g_nf_conn_statis_cache, nf_conn);
		goto out;
	}

	atomic_inc(&g_nf_conn_statis_num);
	hlist_add_head_rcu(&nf_conn->hlnode, &g_nf_conn_statits_list);
out:
	spin_unlock_bh(&g_nf_conn_statis_lock);
}

void nf_conn_statis_try_delete_entry(const struct nf_conn *ct)
{
	struct nf_conn_statis_entry *nf_conn;
	struct nf_conn_statis_key key = {0};

	if (!nf_conn_statis_get_key(ct, &key))
		return;

	spin_lock_bh(&g_nf_conn_statis_lock);
	nf_conn = nf_conn_statis_find(&g_nf_conn_statis_tbl, &key);
	if (!nf_conn)
		goto out;

	if (!atomic_dec_and_test(&nf_conn->use)) {
		if (nf_conn->is_report &&
		    time_after_eq(jiffies, (nf_conn->report_expires + 180 * HZ)))
			nf_conn->is_report = 0;
		goto out;
	}

	if (atomic_dec_and_test(&g_nf_conn_statis_num))
		atomic_set(&g_nf_conn_statis_num, 1);
	hlist_del_init_rcu(&nf_conn->hlnode);
	rhashtable_remove_fast(&g_nf_conn_statis_tbl, &nf_conn->rhnode, nf_conn_rht_params);
	call_rcu(&nf_conn->rcu, nf_conn_statis_rcu_free);
out:
	spin_unlock_bh(&g_nf_conn_statis_lock);
}

static int nf_conn_statis_entry_dump(struct seq_file *file, void *iter)
{
	struct nf_conn_statis_entry *conn = NULL;

	rcu_read_lock();
	hlist_for_each_entry_rcu(conn, &g_nf_conn_statits_list, hlnode) {
		if (conn->key.l3num == NFPROTO_IPV4)
			seq_printf(file, "Ip:%-20pI4 Use:%d UseRate:%d (per second)\n",
				   &conn->key.ipv4, atomic_read(&conn->use),
				   atomic_read(&conn->per_sec_use));
	}

	hlist_for_each_entry_rcu(conn, &g_nf_conn_statits_list, hlnode) {
		if (conn->key.l3num == NFPROTO_IPV6)
			seq_printf(file, "Ip:%-40pI6 Use:%d UseRate:%d (per second)\n",
				   &conn->key.ipv6, atomic_read(&conn->use),
				   atomic_read(&conn->per_sec_use));
	}
	rcu_read_unlock();

	return 0;
}

static int nf_conn_statis_proc_open(struct inode *inode, struct file *file)
{
	return single_open(file, nf_conn_statis_entry_dump, NULL);
}

static struct proc_ops nf_conntrack_ops = {
	.proc_open = nf_conn_statis_proc_open,
	.proc_read = seq_read,
	.proc_release = seq_release,
	.proc_lseek = seq_lseek,
};

int nf_conn_statis_init(void)
{
	int ret;
	struct proc_dir_entry *file;

	g_nf_conn_statis_cache = kmem_cache_create("nf_conn_statis_cache",
						   sizeof(struct nf_conn_statis_entry), 0,
						   SLAB_HWCACHE_ALIGN, NULL);
	if (!g_nf_conn_statis_cache)
		goto cache_create_fail;

	ret = rhashtable_init(&g_nf_conn_statis_tbl, &nf_conn_rht_params);
	if (ret)
		goto rhashtable_fail;

	file = proc_create_data("nf_conntrack_statistics", 0600, init_net.proc_net,
				&nf_conntrack_ops, NULL);
	if (!file)
		goto proc_create_fail;

	atomic_set(&g_nf_conn_statis_num, 1);
	INIT_HLIST_HEAD(&g_nf_conn_statits_list);
	spin_lock_init(&g_nf_conn_statis_lock);
	return 0;

proc_create_fail:
	rhashtable_destroy(&g_nf_conn_statis_tbl);
rhashtable_fail:
	kmem_cache_destroy(g_nf_conn_statis_cache);
cache_create_fail:
	return -1;
}

void nf_conn_statis_fini(void)
{
	remove_proc_entry("nf_conntrack_statistics", init_net.proc_net);
	rhashtable_destroy(&g_nf_conn_statis_tbl);
	kmem_cache_destroy(g_nf_conn_statis_cache);
}

